# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/15 14:32 
@Author  : ZhangShenao 
@File    : 1.使用Embedding模型生成词向量.py
@Desc    : 使用Embedding模型生成词向量
"""
import os

import dotenv
from langchain_openai import OpenAIEmbeddings

from cosine_similarity import cosine_similarity

# 加载环境变量
dotenv.load_dotenv()

# 创建Embedding模型
embeddings = OpenAIEmbeddings(
    model='text-embedding-3-large',  # 选择text-embedding-3-small模型
    dimensions=1024,  # 指定Embedding生成的向量维度
    check_embedding_ctx_length=False,
    openai_api_base=os.getenv('OPENAI_API_BASE')
)

# 使用Embedding模型生成词向量列表
document_vectors = embeddings.embed_documents([
    'zsa是一个喜欢打篮球的人',
    '这个喜欢打篮球的人叫zsa',
    '道阻且长，行则将至',
])

print(len(document_vectors))  # 生成3个词向量
print(len(document_vectors[0]))  # 每个词向量的维度都是100

# vector1与vector2在语义上更相似
# vector1与vector2的余弦相似度: 0.870687996049608
print(f'vector1与vector2的余弦相似度: {cosine_similarity(document_vectors[0], document_vectors[1])}')

# vector1与vector3的余弦相似度: 0.1476896129806336
print(f'vector1与vector3的余弦相似度: {cosine_similarity(document_vectors[0], document_vectors[2])}')
